//
// Sample: Blocking IPv4/IPv6 Server
//
// Files:
//      bserver.cpp     - this file
//      resolve.cpp     - routines for resovling addresses, etc.
//      resolve.h       - header file for resolve.c
// Usage:
//      bserver.exe [options]
//          -a 4|6     Address family, 4 = IPv4, 6 = IPv6 [default = IPv4]
//          -b size    Size of send/recv buffer in bytes
//          -e port    Port number
//          -r port    Port number to relay
//          -l addr    Local address to bind to [default INADDR_ANY for IPv4 or INADDR6_AN
//
#include <winsock2.h>
#include <ws2tcpip.h>
#include <windows.h>
#include <stdio.h>
#include <stdlib.h>

#include "resolve.h"

#define DEFAULT_BUFFER_SIZE     4096    // default buffer size
#define MAX_LISTEN_SOCKETS      8       // maximum listening sockets

int gAddressFamily = AF_UNSPEC,         // default to unspecified    
gBufferSize    = DEFAULT_BUFFER_SIZE;

char *gBindAddr    = NULL,              // local interface to bind to
*gBindPort    = "5150",            // local port to bind to
*gRelayPort   = "8484";

char* REQUEST_CONNECT = "connect to 127.0.0.1:8585",
      *RESPONSE_OK    = "connect ok",
      *RESPONSE_FAIL = "connect fail";

struct _BUFFER_OBJ;

//
// Allocated for each client connection
//
typedef struct _CONNECTION_OBJ
{
    SOCKET      s;              // Client socket
    HANDLE      hRecvSema;      // Semaphore incremented for each receive

    struct _BUFFER_OBJ *PendingSendHead,    // List of pending buffers to send
        *PendingSendTail;    // End of the list

    CRITICAL_SECTION    SendRecvQueueCritSec; // Protect access to this structure

} CONNECTION_OBJ;

//
// Allocate for bouncing two socket
//
typedef struct _BOUNCING_PARAM
{
    CONNECTION_OBJ* ConnObjFrom;
    CONNECTION_OBJ* ConnObjTo;
}BOUNCING_PARAM;

//
// Allocate for bouncing two socket
//
typedef struct _SERVER_PAIRE
{
    SOCKET serverSocket;
    SOCKET relaySocket;
}SERVER_PAIRE;

//
// Allocated for each receiver posted
//    Each receive thread allocates one of these for a receive operation.
//    After data is read, this object is queued for the send thread to
//    echo back to the client (sender).
//
typedef struct _BUFFER_OBJ
{
    char        *buf;           // Data buffer for data
    int          buflen;        // Length of buffer or number of bytes contained in buffer    
    struct _BUFFER_OBJ *next;

} BUFFER_OBJ;

//
// Statistics counters
//
volatile LONG gBytesRead=0,
gBytesSent=0,
gStartTime=0,
gBytesReadLast=0,
gBytesSentLast=0,
gStartTimeLast=0,
gConnectedClients=0;


//
// Function: usage
//
// Description:
//      Prints usage information and exits the process.
//
void usage(char *progname)
{
    fprintf(stderr, "usage: %s [-a 4|6] [-e port] [-r relay-port] [-l local-addr]\n",
        progname);
    fprintf(stderr, "  -a 4|6     Address family, 4 = IPv4, 6 = IPv6 [default = IPv4]\n"
        "  -b size    Buffer size for send/recv [default = %d]\n"
        "  -e port    Port number [default = %s]\n"
        "  -r port    Relay port number [default = %s]\n"
        "  -l addr    Local address to bind to [default INADDR_ANY for IPv4 or INADDR6_ANY for IPv6]\n",                    
        gBufferSize,
        gBindPort,
        gRelayPort
        );
    ExitProcess(-1);
}

//
// Function: GetConnectionObj
//
// Description:
//    This routine allocates a CONNECTION_OBJ structure and initializes its
//    members. For sake of simplicity, this routine allocates a object from
//    the process heap. For better performance, you should modify this to
//    maintain a look aside list of structures already allocated and freed
//    and only allocate from the heap when necessary.
//
CONNECTION_OBJ *GetConnectionObj(SOCKET s)
{
    CONNECTION_OBJ *newobj=NULL;

    // Allocate the object
    newobj = (CONNECTION_OBJ *)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(CONNECTION_OBJ));
    if (newobj == NULL)
    {
        printf("GetConnectionObj: HeapAlloc failed: %d\n", WSAGetLastError());
        ExitProcess(-1);
    }
    newobj->s = s;

    // Create the semaphore for signaling the send thread
    newobj->hRecvSema = CreateSemaphore(NULL, 0, 0x0FFFFFFF, NULL);
    if (newobj->hRecvSema == NULL)
    {
        fprintf(stderr, "GetConnectionObj: CreateSemaphore failed: %d\n", GetLastError());
        ExitProcess(-1);
    }

    InitializeCriticalSection(&newobj->SendRecvQueueCritSec);

    return newobj;
}

//
// Function: FreeConnectionObj
//
// Description:
//    This routine frees a CONNECTIN_OBJ. It first frees the critical section.
//    See the comment for GetConnectionObj about using lookaside lists.
//
void FreeConnectionObj(CONNECTION_OBJ *obj)
{
    DeleteCriticalSection(&obj->SendRecvQueueCritSec);
    HeapFree(GetProcessHeap(), 0, obj);
}

//
// Function: GetBufferObj
// 
// Description:
//    Allocate a BUFFER_OBJ. Each receive posted by a receive thread allocates
//    one of these. After the recv is successful, the BUFFER_OBJ is queued for
//    sending by the send thread. Again, lookaside lists may be used to increase
//    performance.
//
BUFFER_OBJ *GetBufferObj(int buflen)
{
    BUFFER_OBJ *newobj=NULL;

    // Allocate the object
    newobj = (BUFFER_OBJ *)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(BUFFER_OBJ));
    if (newobj == NULL)
    {
        fprintf(stderr, "GetBufferObj: HeapAlloc failed: %d\n", GetLastError());
        ExitProcess(-1);
    }
    // Allocate the buffer
    newobj->buf = (char *)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(BYTE) *buflen);
    if (newobj->buf == NULL)
    {
        fprintf(stderr, "GetBufferObj: HeapAlloc failed: %d\n", GetLastError());
        ExitProcess(-1);
    }
    newobj->buflen = buflen;

    return newobj;
}

//
// Function: FreeBufferObj
// 
// Description:
//    Free the buffer object.
//
void FreeBufferObj(BUFFER_OBJ *obj)
{
    HeapFree(GetProcessHeap(), 0, obj->buf);
    HeapFree(GetProcessHeap(), 0, obj);
}

//
// Function: EnqueueBufferObj
//
// Description:
//   Queue up a receive buffer for this connection. After enqueueing the receiver
//   will release the counting semaphore which will release the sender to 
//   dequeue a buffer and send it.
//
void EnqueueBufferObj(CONNECTION_OBJ *conn, BUFFER_OBJ *obj)
{
    EnterCriticalSection(&conn->SendRecvQueueCritSec);

    if (conn->PendingSendHead == NULL)
    {
        // Queue is empty
        conn->PendingSendHead = conn->PendingSendTail = obj;
    }
    else
    {
        // Put new object at the end 
        conn->PendingSendTail->next = obj;
        conn->PendingSendTail = obj;
    }
    LeaveCriticalSection(&conn->SendRecvQueueCritSec);
}

// 
// Function: DequeueBufferObj
//
// Description:
//    Remove a BUFFER_OBJ from the given connection's queue for sending.
//
BUFFER_OBJ *DequeueBufferObj(CONNECTION_OBJ *conn)
{
    BUFFER_OBJ *ret=NULL;

    EnterCriticalSection(&conn->SendRecvQueueCritSec);

    if (conn->PendingSendTail != NULL)
    {
        // Queue is non empty
        ret = conn->PendingSendHead;

        conn->PendingSendHead = conn->PendingSendHead->next;
        if (conn->PendingSendTail == ret)
        {
            // Item is the only item in the queue
            conn->PendingSendTail = NULL;
        }
    }

    LeaveCriticalSection(&conn->SendRecvQueueCritSec);

    return ret;
}

//
// Function: ValidateArgs
//
// Description:
//      Parses the command line arguments and sets up some global 
//      variables.
//
void ValidateArgs(int argc, char **argv)
{
    int     i;

    for(i=1; i < argc ;i++)
    {
        if (((argv[i][0] != '/') && (argv[i][0] != '-')) || (strlen(argv[i]) < 2))
            usage(argv[0]);
        else
        {
            switch (tolower(argv[i][1]))
            {
            case 'a':               // address family - IPv4 or IPv6
                if (i+1 >= argc)
                    usage(argv[0]);
                if (argv[i+1][0] == '4')
                    gAddressFamily = AF_INET;
                else if (argv[i+1][0] == '6')
                    gAddressFamily = AF_INET6;
                else
                    usage(argv[0]);
                i++;
                break;
            case 'b':               // buffer size for send/recv
                if (i+1 >= argc)
                    usage(argv[0]);
                gBufferSize = atol(argv[++i]);
                break;
            case 'e':               // endpoint - port number
                if (i+1 >= argc)
                    usage(argv[0]);
                gBindPort = argv[++i];
                break;
            case 'r':               // relay - port number
                if (i+1 >= argc)
                    usage(argv[0]);
                gRelayPort = argv[++i];
                break;
            case 'l':               // local address for binding
                if (i+1 >= argc)
                    usage(argv[0]);
                gBindAddr = argv[++i];
                break;                
            default:
                usage(argv[0]);
                break;
            }
        }
    }
}

// 
// Function: SendThread
//
// Description:
//    This is the send thread started for each client connection.
//    This thread waits for the semaphore to be signaled indicating that
//    the receive thread has queued a buffer for sending.
//
DWORD WINAPI SendThread(LPVOID lpParam)
{
    CONNECTION_OBJ *ConnObj=NULL;
    BUFFER_OBJ     *BuffObj=NULL;
    int             rc,
        nleft,
        idx;


    // Retrieve the connection object
    ConnObj = (CONNECTION_OBJ *)lpParam;

    while (1)
    {
        // Wait for the receive thread to signal us
        rc = WaitForSingleObject(ConnObj->hRecvSema, INFINITE);
        if (rc == WAIT_FAILED || rc == WAIT_TIMEOUT)
        {
            fprintf(stderr, "WaitForSingleObject failed: %d\n", GetLastError());
            ExitProcess(-1);
        }

        // Retrieve the first buffer from this connection's queue
        BuffObj = DequeueBufferObj(ConnObj);

        //
        // If the this receive by the receive thread indicated zero bytes then
        // the connection has been gracefully closed. Otherwise, if an error
        // was indicated then the connection was aborted.
        //
        if (((BuffObj->buflen == 0) || (BuffObj->buflen == SOCKET_ERROR)))
        {
            FreeBufferObj(BuffObj);
            BuffObj = NULL;
            break;
        }



        // Otherwise send the buffer on the connection socket
        nleft = BuffObj->buflen;
        idx = 0;

        while (nleft > 0)
        {
            rc = send(
                ConnObj->s,
                &BuffObj->buf[idx],
                nleft,
                0
                );
            if (rc == SOCKET_ERROR)
            {
                break;
            }
            else
            {
                nleft -= rc;
                idx += rc;
            }
        }

        if (rc == SOCKET_ERROR)
        {
            printf("SendThread: send(to) failed: %d\n", WSAGetLastError());
            break;
        }
        else if (rc > 0)
        {
            // Increment the statistics
            InterlockedExchangeAdd(&gBytesSent, rc);
            InterlockedExchangeAdd(&gBytesSentLast, rc);
        }

        FreeBufferObj(BuffObj);
        BuffObj = NULL;
    }

    // Close the connection's socket
    closesocket(ConnObj->s);

    FreeConnectionObj(ConnObj);

    ExitThread(0);
    return 0;
}

//
DWORD WINAPI BouncThread(LPVOID lpParam)
{
    BOUNCING_PARAM *BouncParam = NULL;
    CONNECTION_OBJ *ConnObjFrom=NULL;
    CONNECTION_OBJ *ConnObjTo=NULL;
    BUFFER_OBJ     *BuffObj=NULL;
    int             rc;

    // Retrieve the connection object
    BouncParam = (BOUNCING_PARAM *)lpParam;
    ConnObjFrom = BouncParam->ConnObjFrom;
    ConnObjTo = BouncParam->ConnObjTo;

    // loop until the connection is closed or is aborted/terminated
    while (1)
    {
        // Allocate the buffer for stream send/recv
        BuffObj = GetBufferObj(gBufferSize);

        rc = recv(
            ConnObjFrom->s, 
            BuffObj->buf, 
            BuffObj->buflen, 
            0);
        BuffObj->buflen = rc;

        // Queue the receive buffer for sending and signal the send thread
        EnqueueBufferObj(ConnObjTo, BuffObj);

        ReleaseSemaphore(ConnObjTo->hRecvSema, 1, NULL);

        if (rc == 0 || rc == SOCKET_ERROR)
        {
            break;
        }
        else if (rc != SOCKET_ERROR)
        {
            // Increment the statistics
            InterlockedExchangeAdd(&gBytesRead, rc);
            InterlockedExchangeAdd(&gBytesReadLast, rc);
        }

    }

    ExitThread(0);
    return 0;
}


//
// Function: ServerListenThread
//
// Description:
//    This function is spawned for each listening or receive thread
//    depending on whether the server is started for UDP or TCP. In
//    reality there will only be two server thread, one for IPv4
//    and one for IPv6.
//
DWORD WINAPI ServerListenThread(LPVOID lpParam)
{
    SERVER_PAIRE *ServerPaire=NULL;
    HANDLE          hThread = NULL;
    SOCKET          s;    
    SOCKET          srelay;  
    int             rc;

    ServerPaire = (SERVER_PAIRE*) lpParam;
    s = ServerPaire->serverSocket;
    srelay = ServerPaire->relaySocket;

    SOCKADDR_STORAGE saAccept;      // client address
    SOCKET           ns;            // client socket
    int              acceptlen = sizeof(SOCKADDR_STORAGE);

    SOCKADDR_STORAGE saAcceptRelay;      // relay client address
    SOCKET           nsRelay;            // relay client socket    

    // Trans listen
    rc = listen(s, 200);
    if (rc == SOCKET_ERROR)
    {
        fprintf(stderr, "listen failed: %d\n", WSAGetLastError());
        ExitThread(-1);
    }else{
        fprintf(stdout,"server start listen\n");
    }

    // Relay listen
    rc = listen(srelay, 200);
    if (rc == SOCKET_ERROR)
    {
        fprintf(stderr, "listen failed: %d\n", WSAGetLastError());
        ExitThread(-1);
    }else{
        fprintf(stdout,"relay start listen\n");
    }

    while (TRUE)
    {
        printf("wait and accept client connection...\n");
        ns = accept(s, (SOCKADDR *)&saAccept, &acceptlen);
        if (ns == INVALID_SOCKET)
        {
            fprintf(stderr, "accept failed: %d\n", WSAGetLastError());
            return -1;
        }else{            
            printf("Accepted connection from: ");
            PrintAddress((SOCKADDR *)&saAccept, acceptlen); 
            printf("\n");            
        }

        InterlockedIncrement(&gConnectedClients);        

        bool bOneMoreRelay = false;
        char* buff = new char[256];

        while(!bOneMoreRelay){            
            printf("wait and accept relay connection...\n");
            nsRelay = accept(srelay, (SOCKADDR *)&saAcceptRelay, &acceptlen);
            if (nsRelay == INVALID_SOCKET)
            {
                fprintf(stderr, "accept failed: %d\n", WSAGetLastError());
                break;
            }else{            
                printf("accepted relay connection from: ");
                PrintAddress((SOCKADDR *)&saAcceptRelay, acceptlen); 
                printf("\n");            
            }

            printf("send request to client...\n");
            rc = send(ns,REQUEST_CONNECT,strlen(REQUEST_CONNECT),NULL);
            if(rc == SOCKET_ERROR){
                fprintf(stderr, "send failed: %d\n", WSAGetLastError());
                break;
            }else{
                printf("request sent\n");
            }

            printf("receive request response...\n");
            rc = recv(ns,buff,256,NULL);
            if(rc == SOCKET_ERROR){
                fprintf(stderr, "receive failed: %d\n", WSAGetLastError());
                break;
            }else{
                printf("received response\n");
            }
            
            printf("response ok, create send thread\n");
            CONNECTION_OBJ* ConnObjFrom = GetConnectionObj(ns);
            CONNECTION_OBJ* ConnObjTo = GetConnectionObj(nsRelay);
            HANDLE hThread = CreateThread(NULL, 0, SendThread, (LPVOID)ConnObjFrom, 0, NULL);
            if (hThread == NULL)
            {
                fprintf(stderr, "CreateThread failed: %d\n", GetLastError());
                ExitThread(-1);
            }
            CloseHandle(hThread);

            hThread = CreateThread(NULL, 0, SendThread, (LPVOID)ConnObjTo, 0, NULL);
            if (hThread == NULL)
            {
                fprintf(stderr, "CreateThread failed: %d\n", GetLastError());
                ExitThread(-1);
            }
            CloseHandle(hThread);

            printf("create bouncing thread\n");
            BOUNCING_PARAM* pBnc = new BOUNCING_PARAM;
            pBnc->ConnObjFrom = ConnObjFrom;
            pBnc->ConnObjTo = ConnObjTo;

            hThread = CreateThread(NULL, 0, BouncThread, (LPVOID)pBnc, 0, NULL);
            if (hThread == NULL)
            {
                fprintf(stderr, "CreateThread failed: %d\n", GetLastError());
                ExitThread(-1);
            }
            CloseHandle(hThread);

            BOUNCING_PARAM* pBnc2 = new BOUNCING_PARAM;
            pBnc2->ConnObjFrom = ConnObjTo;
            pBnc2->ConnObjTo = ConnObjFrom;

            hThread = CreateThread(NULL, 0, BouncThread, (LPVOID)pBnc2, 0, NULL);
            if (hThread == NULL)
            {
                fprintf(stderr, "CreateThread failed: %d\n", GetLastError());
                ExitThread(-1);
            }
            CloseHandle(hThread);
            bOneMoreRelay = true;
        }
    }

    closesocket(srelay);
    closesocket(s);    

    ExitThread(0);
    return 0;
}


//
// Function: main
//
// Description:
//      This is the main program. It parses the command line and creates
//      the main socket. For UDP this socket is used to receive datagrams.
//      For TCP the socket is used to accept incoming client connections.
//      Each client TCP connection is handed off to a worker thread which
//      will receive any data on that connection until the connection is
//      closed.
//
int __cdecl main(int argc, char **argv)
{
    WSADATA          wsd;
    SOCKET           s;
    SOCKET           srelay;
    HANDLE           thread;
    int              rc;
    struct addrinfo *res=NULL;        

    ValidateArgs(argc, argv);

    if (WSAStartup(MAKEWORD(2,2), &wsd) != 0)
    {
        fprintf_s(stderr, "unable to load Winsock!\n");
        return -1;
    }

    //////////////////////////////////////////////////////////////////////////
    /*Create server socket and bind it*/
    printf("Local address: %s; Port: %s; Family: %d\n",
        gBindAddr, gBindPort, gAddressFamily);

    res = ResolveAddress(gBindAddr, gBindPort, gAddressFamily, SOCK_STREAM, IPPROTO_TCP);
    if (res == NULL)
    {
        fprintf(stderr, "ResolveAddress failed to return any addresses!\n");
        return -1;
    }
    else
    {
        PrintAddress(res->ai_addr, res->ai_addrlen); printf("\n");

        // create the socket
        s = socket(res->ai_family, res->ai_socktype, res->ai_protocol);
        if (s == INVALID_SOCKET)
        {
            fprintf(stderr,"socket failed: %d\n", WSAGetLastError());
            return -1;
        }

        // bind the socket to a local address and port
        rc = bind(s, res->ai_addr, res->ai_addrlen);
        if (rc == SOCKET_ERROR)
        {
            fprintf(stderr, "bind failed: %d\n", WSAGetLastError());
            return -1;
        }
    }

    freeaddrinfo(res);
    res = NULL;

    //////////////////////////////////////////////////////////////////////////
    /*Create relay socket and bind it*/
    printf("Relay address: %s; Port: %s; Family: %d\n",
        gBindAddr, gRelayPort, gAddressFamily);

    res = ResolveAddress(gBindAddr, gRelayPort, gAddressFamily, SOCK_STREAM, IPPROTO_TCP);
    if (res == NULL)
    {
        fprintf(stderr, "ResolveAddress failed to return any addresses!\n");
        return -1;
    }else
    {
        PrintAddress(res->ai_addr, res->ai_addrlen); printf("\n");

        // create the socket
        srelay = socket(res->ai_family, res->ai_socktype, res->ai_protocol);
        if (srelay == INVALID_SOCKET)
        {
            fprintf(stderr,"socket failed: %d\n", WSAGetLastError());
            return -1;
        }

        // bind the socket to a local address and port
        rc = bind(srelay, res->ai_addr, res->ai_addrlen);
        if (rc == SOCKET_ERROR)
        {
            fprintf(stderr, "bind failed: %d\n", WSAGetLastError());
            return -1;
        }

    }
    freeaddrinfo(res);
    res = NULL;

    //////////////////////////////////////////////////////////////////////////
    /* Create a listen thread*/
    SERVER_PAIRE* pServerPaire = new SERVER_PAIRE;
    pServerPaire->serverSocket = s;
    pServerPaire->relaySocket = srelay;

    thread = CreateThread(
        NULL, 
        0, 
        ServerListenThread, 
        (LPVOID)pServerPaire,
        0,
        NULL
        );

    if (thread == NULL)
    {
        fprintf(stderr, "CreateThread failed: %d\n", GetLastError());
        return -1;
    }


    /*Now system started*/
    gStartTime = gStartTimeLast = GetTickCount();

    while (1)
    {
        // Wait for the listening thread to exit, also print out statistics
        // when the wait times out
        rc = WaitForSingleObject(thread, 30000);
        if (rc == WAIT_FAILED)
        {
            fprintf(stderr, "WaitForSingleObject failed: %d\n", GetLastError());
            return -1;
        }
        else if (rc == WAIT_TIMEOUT)
        {
            ULONG       bps, tick, elapsed;

            tick = GetTickCount();

            elapsed = (tick - gStartTime) / 1000;

            printf("\n");

            bps = gBytesSent / elapsed;
            printf("Average BPS sent: %lu [%lu]\n", bps, gBytesSent);

            bps = gBytesRead / elapsed;
            printf("Average BPS read: %lu [%lu]\n", bps, gBytesRead);

            elapsed = (tick - gStartTimeLast) / 1000;

            bps = gBytesSentLast / elapsed;
            printf("Current BPS sent: %lu\n", bps);

            bps = gBytesReadLast / elapsed;
            printf("Current BPS read: %lu\n", bps);

            InterlockedExchange(&gBytesSentLast, 0);
            InterlockedExchange(&gBytesReadLast, 0);

            printf("Current Connections: %lu\n", gConnectedClients);

            gStartTimeLast = tick;
        }
        else
        {
            break;
        }
    }

    // Close the thread handles opened
    CloseHandle(thread);

    WSACleanup();
    return 0;
}
